import pickle
from io import BytesIO

import tensorflow as tf  # 导入tensorflow库

from security.api.symmetric import encrypt_stream, decrypt_stream


# 加载数据
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 搭建模型
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax))
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 数据预处理
x_train, y_train = x_train[:1000] / 255, y_train[:1000]
x_test, y_test = x_test[:1000] / 255, y_test[:1000]

# 训练
model.fit(x_train, y_train, epochs=5)

# 评估
val_loss, val_acc = model.evaluate(x_test, y_test)
print('first evaluate loss: {} acc: {}'.format(val_loss, val_acc))

# 加密保存模型
key = "1c244fb235498d976166fd381368c60707149a9527559c5ced4aed36cd78de3b"

# 加密保存模型结构
f_obj = BytesIO()
f_obj.write(model.to_json().encode('utf-8'))
with open("model_struct.json", 'wb') as f:
    for i in encrypt_stream(key, f_obj):
        f.write(i)

# 加密保存模型权重
f_obj = BytesIO()
weights = model.get_weights()
f_obj.write(pickle.dumps(weights))
with open("model_weights.pkl", 'wb') as f:
    for i in encrypt_stream(key, f_obj):
        f.write(i)


# 解密模型
f_obj = BytesIO()
for i in decrypt_stream(key, 'model_struct.json'):
    f_obj.write(i)

# 加载模型结构
new_model = tf.keras.models.model_from_json(f_obj.getvalue())
f_obj.close()

# 解密并加载模型权重
f_obj = BytesIO()
for i in decrypt_stream(key, 'model_weights.pkl'):
    f_obj.write(i)
new_model.set_weights(pickle.loads(f_obj.getvalue()))
f_obj.close()

# 推理
res = new_model.predict(x_test[:10])
print('predict', ' '.join([str(tf.argmax(i).numpy()) for i in res]))
print('real  ', y_test[:10])

# 模型编译，用于评估或二次训练，不需要评估或二次训练可以不编译
new_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 评估
val_loss, val_acc = new_model.evaluate(x_test, y_test)

print('second evaluate loss: {} acc: {}'.format(val_loss, val_acc))

# 保存成ckpt和pb格式(明文)
new_model.save('ckpt_dir')
